# aruco_detector.py

import cv2
import numpy as np
from utils import frame_to_bgr_image
from camera import Camera
import time

ESC_KEY = 27

# Define thresholds
T_VEC_THRESHOLD = 5
EULER_ANGLE_THRESHOLD = 5.0


class DetectionResult:
    def __init__(self, ids=None, euler_angles=None, tvecs=None):
        self.ids = ids
        self.euler_angles = euler_angles
        self.tvecs = tvecs


class ArUcoDetector:
    def __init__(self):
        self.camera_matrix, self.dist_coeffs = self.read_camera_parameters()
        self.aruco_dict = cv2.aruco.Dictionary_get(cv2.aruco.DICT_5X5_250)
        self.parameters = cv2.aruco.DetectorParameters_create()
        self.prev_tvecs = None
        self.prev_euler_angles = None

    def read_camera_parameters(self):
        camera_matrix = np.array(
            [
                [303.155313, 0.0, 311.192343],
                [0.0, 306.224939, 200.859069],
                [0.0, 0.0, 1.0],
            ],
            dtype=np.float32,
        )
        dist_coeffs = np.array(
            [-0.003614, 0.012834, -0.002176, -0.003698, 0.0], dtype=np.float32
        ).reshape(5, 1)
        return camera_matrix, dist_coeffs

    def rotation_matrix_to_euler_angles(self, R):
        assert R.shape == (3, 3)
        sy = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
        singular = sy < 1e-6
        if not singular:
            x = np.arctan2(R[2, 1], R[2, 2])
            y = np.arctan2(-R[2, 0], sy)
            z = np.arctan2(R[1, 0], R[0, 0])
        else:
            x = np.arctan2(-R[1, 2], R[1, 1])
            y = np.arctan2(-R[2, 0], sy)
            z = 0
        # Convert radians to degrees
        x_deg = np.degrees(x)
        y_deg = np.degrees(y)
        z_deg = np.degrees(z)

        return np.array([x_deg, y_deg, z_deg])

    def euler_angles_to_rotation_matrix(self, euler_angles):
        # Convert degrees to radians
        euler_angles = np.radians(euler_angles)

        # Create individual rotation matrices
        R_x = np.array(
            [
                [1, 0, 0],
                [0, np.cos(euler_angles[0]), -np.sin(euler_angles[0])],
                [0, np.sin(euler_angles[0]), np.cos(euler_angles[0])],
            ]
        )

        R_y = np.array(
            [
                [np.cos(euler_angles[1]), 0, np.sin(euler_angles[1])],
                [0, 1, 0],
                [-np.sin(euler_angles[1]), 0, np.cos(euler_angles[1])],
            ]
        )

        R_z = np.array(
            [
                [np.cos(euler_angles[2]), -np.sin(euler_angles[2]), 0],
                [np.sin(euler_angles[2]), np.cos(euler_angles[2]), 0],
                [0, 0, 1],
            ]
        )

        # Combine rotations
        R = np.dot(R_z, np.dot(R_y, R_x))

        return R

    def detect_aruco(self, color_image):
        gray = cv2.cvtColor(color_image, cv2.COLOR_BGR2GRAY)
        corners, ids, rejected = cv2.aruco.detectMarkers(
            gray, self.aruco_dict, parameters=self.parameters
        )

        euler_angles, tvecs = None, None
        if ids is not None:
            rvecs, tvecs, _ = cv2.aruco.estimatePoseSingleMarkers(
                corners, 99.5, self.camera_matrix, self.dist_coeffs
            )

            euler_angles = []
            valid_detection = False
            for i in range(len(ids)):
                R, _ = cv2.Rodrigues(rvecs[i])
                current_euler_angles = self.rotation_matrix_to_euler_angles(R)

                if (self.prev_tvecs is None or self.prev_euler_angles is None) or (
                    np.linalg.norm(tvecs[i] - self.prev_tvecs[i]) > T_VEC_THRESHOLD
                    or np.any(
                        np.abs(current_euler_angles - self.prev_euler_angles[i])
                        > EULER_ANGLE_THRESHOLD
                    )
                ):
                    valid_detection = True
                    euler_angles.append(current_euler_angles)
                else:
                    euler_angles.append(self.prev_euler_angles[i])
                    tvecs[i] = self.prev_tvecs[i]

            if valid_detection:
                self.prev_tvecs = tvecs
                self.prev_euler_angles = euler_angles
            color_image = cv2.aruco.drawDetectedMarkers(color_image, corners, ids)
            for i in range(len(ids)):
                R = self.euler_angles_to_rotation_matrix(euler_angles[i])
                rvec, _ = cv2.Rodrigues(R)
                color_image = cv2.aruco.drawAxis(
                    color_image,
                    self.camera_matrix,
                    self.dist_coeffs,
                    rvec,
                    tvecs[i],
                    110,
                )

        return color_image, DetectionResult(ids, euler_angles, tvecs)


if __name__ == "__main__":
    camera = Camera()
    if not camera.configure_pipeline():
        exit()
    if not camera.start_pipeline():
        exit()

    detector = ArUcoDetector()

    start_time = time.time()
    frame_count = 0
    try:
        while True:
            color_image = camera.get_frame()
            if color_image is None:
                continue

            detected_image, detection_result = detector.detect_aruco(color_image)
            cv2.imshow("Color Viewer", detected_image)
            if detection_result.ids is not None:
                print(f"Detected IDs: {detection_result.ids}")
                print(f"Euler Angles: {detection_result.euler_angles}")
                print(f"Translation Vectors: {detection_result.tvecs}")

            frame_count += 1
            elapsed_time = time.time() - start_time
            if elapsed_time > 1.0:
                fps = frame_count / elapsed_time
                print(f"FPS: {fps:.2f}")
                start_time = time.time()
                frame_count = 0

            key = cv2.waitKey(1)
            if key == ord("q") or key == ESC_KEY:
                break
    except KeyboardInterrupt:
        pass
    finally:
        camera.stop_pipeline()
        cv2.destroyAllWindows()
